import os
import logging
import uuid
from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn
from typing import Optional, List, Dict
import shutil

from vectors import EmbeddingsManager, SUPPORTED_EXTENSIONS
from chatbot import ChatbotManager

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(
    title="Document RAG API",
    description="API for RAG-based document chatbot",
    version="1.0.0"
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

# Global variables
temp_pdf_dir = "temp_pdfs"
os.makedirs(temp_pdf_dir, exist_ok=True)

# Job status tracking
job_status: Dict[str, Dict] = {}

# Initialize managers
embeddings_manager = EmbeddingsManager(
    model_name="moka-ai/m3e-base",
    device="cpu",
    encode_kwargs={"normalize_embeddings": True},
    collection_name="vector_db"
)

chatbot_manager = ChatbotManager(
    model_name="moka-ai/m3e-base",
    device="cpu",
    encode_kwargs={"normalize_embeddings": True},
    llm_model="glm-4",
    llm_temperature=0.7,
    collection_name="vector_db"
)

# Request/Response models
class MessageRequest(BaseModel):
    query: str

class MessageResponse(BaseModel):
    answer: str

class JobStatusResponse(BaseModel):
    job_id: str
    status: str
    message: str
    filename: Optional[str] = None

@app.get("/")
async def root():
    return {"message": "Welcome to Document RAG API"}

@app.post("/clear-database", response_model=dict)
async def clear_database():
    """Clear the current vector database"""
    try:
        result = embeddings_manager.clear_database()
        return {"status": "success", "message": result}
    except Exception as e:
        logger.error(f"Error clearing database: {e}")
        raise HTTPException(status_code=500, detail=f"Failed to clear database: {str(e)}")

@app.post("/upload-document", response_model=JobStatusResponse)
async def upload_document(background_tasks: BackgroundTasks, file: UploadFile = File(...)):
    """Upload a document and create embeddings"""
    # 获取文件扩展名并检查是否支持
    file_ext = os.path.splitext(file.filename)[-1].lower()
    if file_ext not in SUPPORTED_EXTENSIONS:
        raise HTTPException(
            status_code=400, 
            detail=f"Unsupported file type: {file_ext}. Supported types are: {', '.join(SUPPORTED_EXTENSIONS)}"
        )
    
    try:
        # Generate a job ID
        job_id = str(uuid.uuid4())
        
        # Save the file
        file_path = os.path.join(temp_pdf_dir, file.filename)
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        
        # Initialize job status
        job_status[job_id] = {
            "status": "processing",
            "message": "Document uploaded and processing started",
            "filename": file.filename
        }
        
        # Process embeddings in the background
        background_tasks.add_task(process_document, file_path, job_id)
        
        return JobStatusResponse(
            job_id=job_id,
            status="processing", 
            message="Document uploaded and processing started", 
            filename=file.filename
        )
    except Exception as e:
        logger.error(f"Error uploading document: {e}")
        raise HTTPException(status_code=500, detail=f"Failed to upload document: {str(e)}")

@app.get("/job-status/{job_id}", response_model=JobStatusResponse)
async def get_job_status(job_id: str):
    """Get the status of a document processing job"""
    if job_id not in job_status:
        raise HTTPException(status_code=404, detail="Job not found")
    
    job_info = job_status[job_id]
    return JobStatusResponse(
        job_id=job_id,
        status=job_info["status"],
        message=job_info["message"],
        filename=job_info.get("filename")
    )

@app.post("/chat", response_model=MessageResponse)
async def chat(request: MessageRequest):
    """Get a response from the RAG chatbot"""
    try:
        # Get response from the chatbot
        answer = chatbot_manager.get_response(request.query)
        return {"answer": answer}
    except Exception as e:
        logger.error(f"Error generating response: {e}")
        raise HTTPException(status_code=500, detail=f"Failed to generate response: {str(e)}")

async def process_document(file_path: str, job_id: str):
    """Process a document to create and store embeddings"""
    try:
        logger.info(f"Processing document: {file_path}")
        result = embeddings_manager.create_embeddings(file_path)
        logger.info(f"Document processing completed: {result}")
        
        # Update job status to completed
        job_status[job_id] = {
            "status": "completed",
            "message": "Document processing successfully completed",
            "filename": os.path.basename(file_path),
            "result": result
        }
    except Exception as e:
        error_msg = f"Error processing document: {e}"
        logger.error(error_msg)
        
        # Update job status to failed
        job_status[job_id] = {
            "status": "failed",
            "message": error_msg,
            "filename": os.path.basename(file_path)
        }
    finally:
        # Clean up the file after processing
        if os.path.exists(file_path):
            os.remove(file_path)
            logger.info(f"Removed temporary file: {file_path}")

if __name__ == "__main__":
    uvicorn.run("api_server:app", host="0.0.0.0", port=9888, reload=False) 